import cv2
import numpy as np
from matplotlib import pyplot as plt
import math

img = cv2.imread('../../../../../large_data/pic/sudoku.png',0)

# ======== cnn

kernel_m = np.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
kernel_x = np.array([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]])
kernel_y = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])

def cal_convoluation_size(input, kernel, padding=0, stride=1, dilation=1):
    new_kernel = dilation * (kernel - 1) + 1  # 空洞卷积，空洞数为0时dilation=1
    # 根据公式计算输出，并返回
    return math.floor((input + 2 * padding - new_kernel) / stride + 1)

# 简单版本的直接卷积法：不考虑padding，dilation=1，padding=0
def convoluation(image, kernel):
    image_height, image_width = image.shape
    kernel_height, kernel_width = kernel.shape
    # 计算输出的形状大小
    out_height = cal_convoluation_size(image_height, kernel_height)
    out_width = cal_convoluation_size(image_width, kernel_width)
    output = np.zeros((out_height, out_width))

    # 计算output的每个像素值
    # 先找到目标图(dx, dy)对应原图中的中心点位置(cx, cy)，然后计算
    for dy in range(out_height):
        for dx in range(out_width):
            # 遍历kernel计算输出(output[dy, dx])的像素值
            for ky in range(kernel_height):
                for kx in range(kernel_width):
                    kernel_value = kernel[ky, kx]
                    pixel_value = image[dy + ky, dx + kx]
                    output[dy, dx] += kernel_value * pixel_value

    return output

if __name__ == '__main__':
    res_m = convoluation(img, kernel_m)
    res_x = convoluation(img, kernel_x)
    res_y = convoluation(img, kernel_y)

    images = [img, res_m, res_x, res_y]

    titles = ['Original', 'mean_filter', 'x_gradient ', 'y_gradient']

    for i in range(4):
        plt.subplot(2, 2, i + 1)
        plt.imshow(images[i], 'gray')
        plt.title(titles[i])
        plt.xticks([]), plt.yticks([])
    plt.show()
